#include"spsolver.h"
#include<iostream>
#include<fstream>
#include<cmath>
#include<numeric>
#include<iomanip>
using std::cout;
using std::endl;
using std::ifstream;
using std::accumulate;
using std::setw;


void SPSolver::init()
{
  SystemHF.setup();
  //adc2.setup();
  //  adc2.solveDysonMatrix_Lanczos(3,3);
  
  //setup Channels
  Channels.clear();
  //  read contour from file
  // ifstream fin("1.contour");
  // int l,jj,tz;
  // double kx,ky;
  // complexd k[100];
  // int nums[100];
  // int n;
  // while(true)
  //   {
  //     fin>>l>>jj>>tz;
  //     if(tz==0) break;
  //     int point=0;
  //     int totalnum=0;
  //     while(true)
  // 	{
  // 	  fin>>kx>>ky>>n;
  // 	  k[point]=complexd(kx,ky);
  // 	  if(n==0) break;
  // 	  nums[point]=n;
  // 	  ++point;
  // 	  totalnum+=n;
  // 	}
  //     Channels.push_back( SPChannel(l,jj,tz,totalnum) );
  //     int numindex=0;
  //     for(int i=0;i<point;i++)
  // 	{
  // 	  vector<double> x(nums[i]),w(nums[i]);
  // 	  gauleg(0,1,x,w);
  // 	  for(int j=0;j<nums[i];j++)
  // 	    {
  // 	      Channels[Channels.size()-1].k[numindex+j]=k[i]+(k[i+1]-k[i])*x[j];
  // 	      Channels[Channels.size()-1].wk[numindex+j]=w[j]*(k[i+1]-k[i]);
  // 	    }
  // 	  numindex+=nums[i];
  // 	}
  //   }


  //read rotation contour from file
  ifstream fin("2.contour");
  int l,jj,tz;
  double theta;
  double thetas[100];
  double kmaxs[100];
  double kmax;
  int num1s[100];
  int num2s[100];
  int num1,num2;
  int n;
  while(true)
    {
      fin>>l>>jj>>tz>>kmax>>theta>>num1>>num2;
      if(tz==0) break;
      Channels.push_back( SPChannel(l,jj,tz,num1+num2) );
      double THETA=-theta*Pi/180;
      complexd K=kmax*complexd( cos(-THETA),sin(-THETA) );
      vector<double> x(num1),w(num1);
      gauleg(0,1,x,w);
      for(int i=0;i<num1;i++)
  	{
  	  Channels[Channels.size()-1].k[i]=K*x[i];
  	  Channels[Channels.size()-1].wk[i]=K*w[i];
  	}
      x.resize(num2);
      w.resize(num2);
      gauleg(0,1,x,w);
      for(int i=0;i<num2;i++)
  	{
  	  Channels[Channels.size()-1].k[num1+i]=kmax*complexd( cos(-x[i]*THETA),sin(-x[i]*THETA) );
  	  Channels[Channels.size()-1].wk[num1+i]=kmax*complexd( cos(-x[i]*THETA),sin(-x[i]*THETA) )*w[i]*complexd(0,THETA);
  	}
    }

  
  //debug
  // for(int i=0;i<Channels.size();i++)
  //   {
  //     cout<<Channels[i].l<<"\t"<<Channels[i].jj<<"\t"<<Channels[i].tz<<endl;
  //     for(int j=0;j<Channels[i].k.size();j++)
  // 	{
  // 	  cout<<Channels[i].k[j]<<"\t"<<Channels[i].wk[j]<<endl;
  // 	}
  //     cout<<endl<<endl;
  //   }
}


KDataType SPSolver::hf_k(int i,KDataType k)
{
  KDataType hfk=KDataType();
  int GroupIndex=hfsolver.HFOrbitals[i].GroupIndex;
  int dim=pSystem_Tz->Groups[GroupIndex].size();
  int l=hfsolver.HFOrbitals[i].l;
  double b=pSystem_Tz->lengthb;
  for(int j=0;j<dim;j++)
    {
      int brak=pSystem_Tz->Groups[GroupIndex][j];
      int n=pSystem_Tz->Orbitals[brak].n;
      hfk +=  hfsolver.HFOrbitals[i].Coeff[j] * cmplx_ho_k(n,l,b,k);
    }
  return hfk;
}

KDataType SPSolver::getHFk1k2(KDataType k1,KDataType k2,int group)
{
  KDataType valk1k2=0;
  int bra=pSystem_Tz->Groups[group][0];
    
  int jja=pSystem_Tz->Orbitals[bra].jj;
  int tza=pSystem_Tz->Orbitals[bra].tz;
  int la=pSystem_Tz->Orbitals[bra].l;

  for(int i=0;i<hfsolver.HFOrbitals.size();i++)
    {
      if(hfsolver.HFOrbitals[i].state) continue;//skip particles
      int jjh=hfsolver.HFOrbitals[i].jj;
      int Jmin=abs(jja - jjh)/2;
      int Jmax=(jja + jjh)/2;
      if(Jmin>Jmax) continue;
      int GroupIndex=hfsolver.HFOrbitals[i].GroupIndex;
      int dim=pSystem_Tz->Groups[GroupIndex].size();

      //k1 \dot k2
      if(hfsolver.HFOrbitals[i].tz!=tza) continue;
      if((hfsolver.HFOrbitals[i].l+la+1)%2) continue;//parity
      for(int J=Jmin;J<=Jmax;J++)
	{
	  valk1k2-=(2*J+1)*(jjh+1) * pow(gsl_sf_coupling_3j(jja,2,jjh,1,0,-1),2) * gsl_sf_coupling_6j(jja,jjh,J*2,jja,jjh,2) * k1 * k2 * hf_k(i,k1) * hf_k(i,k2);
	}

    }
  return valk1k2;
}



KDataType SPSolver::getVkkfromHF(int channel,int a,int b)
{
  KDataType Vkk=KDataType();
  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  KDataType k1=Channels[channel].k[a];
  KDataType k2=Channels[channel].k[b];
  KDataType wk1=Channels[channel].wk[a];
  KDataType wk2=Channels[channel].wk[b];
  
  double lengthb=pSystem_Tz->lengthb;
  int totalGroups=pSystem_Tz->Groups.size();
  int group=0;//find the group of l,jj,tz
  for(group=0;group<totalGroups;group++)
    {
      const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][0] ];
      if(HOOrbital.l==l && HOOrbital.jj==jj && HOOrbital.tz==tz)
	break;
    }
  if(group>=totalGroups) return Vkk;
  int dim=pSystem_Tz->Groups[group].size();
  for(int i=0;i<dim;i++)
    for(int j=0;j<dim;j++)
      {
	const System_Tz_HO::OrbitalType & HOOrbital1=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][i] ];
	const System_Tz_HO::OrbitalType & HOOrbital2=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][j] ];
	int n1=HOOrbital1.n;
	int n2=HOOrbital2.n;
	//	Vkk+=hfsolver.HFhamil(group,i,j)*cmplx_ho_k(n1,l,lengthb,k1) * cmplx_ho_k(n2,l,lengthb,k2) * k1 * k2 * sqrt(wk1*wk2);
	int bra=pSystem_Tz->Groups[group][i];
	int ket=pSystem_Tz->Groups[group][j];
	
        Vkk+=hfsolver.SHFDiagram(bra,ket)*cmplx_ho_k(n1,l,lengthb,k1)*cmplx_ho_k(n2,l,lengthb,k2)* k1 * k2 * sqrt(wk1*wk2);
	//	Vkk+=hfsolver.getVHF(bra,ket)*cmplx_ho_k(n1,l,lengthb,k1)*cmplx_ho_k(n2,l,lengthb,k2)* k1 * k2 * sqrt(wk1*wk2);
      }
  //  Vkk-=hbar_c*hbar_c/(A*mc2)*getHFk1k2(k1,k2,group) * k1 * k2 * sqrt(wk1*wk2);
  //  if(a==b) Vkk+=(1-1./A)*hbar_c*hbar_c*0.5/mc2*k1*k1;
  if(abs(Vkk)<1e-8) Vkk=0.;
  return Vkk;
}

double SPSolver::getVrr(int channel,double r1,double r2)
{
  return getVrrfromVkk(channel,r1,r2).real();
}
KDataType SPSolver::getVkkfromVrr(int channel,int a,int b)
{
  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  KDataType k1=Channels[channel].k[a];
  KDataType k2=Channels[channel].k[b];
  KDataType wk1=Channels[channel].wk[a];
  KDataType wk2=Channels[channel].wk[b];

  KDataType Vkk=0.;
  for(int i=0;i<rdim;i++)
    for(int j=0;j<rdim;j++)
      {
  	double r1=rx[i];
  	double wr1=rw[i];
  	double r2=rx[j];
  	double wr2=rw[j];
  	Vkk+=csphbessel_jn(l,k1*r1)*csphbessel_jn(l,k2*r2)*getVrr(channel,r1,r2)*r1*r1*r2*r2*wr1*wr2;
      }
  Vkk*=2./Pi*k1*k2*sqrt(wk1*wk2);
  return Vkk;
}

KDataType SPSolver::getVrrfromVkk(int channel,double r1,double r2)
{
  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  KDataType Vrr=0.;
  int kdim=Channels[channel].numPoints;
  for(int i=0;i<kdim;i++)
    for(int j=0;j<kdim;j++)
      {
	KDataType k1=Channels[channel].k[i];
	KDataType k2=Channels[channel].k[j];
	KDataType wk1=Channels[channel].wk[i];
	KDataType wk2=Channels[channel].wk[j];
	Vrr+=csphbessel_jn(l,k1*r1)*csphbessel_jn(l,k2*r2)*k1*k2*sqrt(wk1*wk2)*getVkkfromHF(channel,i,j);
      }
  Vrr*=2./Pi;
  return Vrr;
}
void SPSolver::printVrr(int channel)
{
  for(int i=0;i<rdim;i++)
    {
      for(int j=0;j<rdim;j++)
	{
	  double r1=rx[i];
	  double r2=rx[j];
	  cout<<i<<"\t"<<j<<"\t"<<r1<<"\t"<<r2<<"\t"<<getVrrfromVkk(channel,r1,r2).real()<<endl;
	}
      cout<<endl;
    }
}
void SPSolver::setVkkMatfromGF(int channel,KMat &h)
{
//  int dimk=Channels[channel].numPoints;
//  h.resize(dimk,dimk);
//  h.setZero();
//  for(int i=0;i<dimk;i++)
//    {
//      KDataType k1=Channels[channel].k[i];
//      h(i,i)=(1-1./A)*hbar_c*hbar_c*0.5/mc2*k1*k1;
//    }
//  int l=Channels[channel].l;
//  int jj=Channels[channel].jj;
//  int tz=Channels[channel].tz;
//  double lengthb=pSystem_Tz->lengthb;
//  int totalBlocks=adc2.blocks.size();
//  int block=0;//find the block of l,jj,tz
//  for(block=0;block<totalBlocks;block++)
//    {
//      if(adc2.blocks[block].l==l && adc2.blocks[block].jj==jj && adc2.blocks[block].tz==tz)
//	break;
//    }
//  if(block>=totalBlocks) return;
//  int dimhf=adc2.blocks[block].dim;
//  KMat T(dimk,dimhf);
//  KMat UHF(dimhf,dimhf);
//  UHF.setZero();
//  MatrixXd U(dimhf,dimhf);
//  for(int i=0,j=0;i<adc2.poles[block].rows();i++)
//    {
//      if(adc2.Z[block].col(i).norm()>0.5)
//	{
//	  U(j,j)=adc2.poles[block][i];
//	  j++;
//	}
//    }
//  for(int i=0;i<dimhf;i++)
//    {
//      for(int j=i;j<dimhf;j++)
//	{
//	  int orb1=SystemHF.Groups[block][i];
//      	  int orb2=SystemHF.Groups[block][j];
//	  UHF(i,j)=U(i,j)-SystemHF.get1B(orb1,orb2);
//	  UHF(j,i)=UHF(i,j);
//	}
//    }
//  for(int i=0;i<dimk;i++)
//    for(int j=0;j<dimhf;j++)
//      {
//	int orb=SystemHF.Groups[block][j];
//	KDataType k1=Channels[channel].k[i];
//	KDataType wk1=Channels[channel].wk[i];
//	T(i,j)=	hf_k(orb,k1) * k1 * sqrt(wk1);
//      }
//  h+=T*UHF*T.transpose();
}
void SPSolver::setVkkMatfromHF(int channel,KMat &h)
{
  int dimk=Channels[channel].numPoints;
  h.resize(dimk,dimk);
  h.setZero();
  for(int i=0;i<dimk;i++)
    {
      KDataType k1=Channels[channel].k[i];
      h(i,i)=(1-1./A)*hbar_c*hbar_c*0.5/mc2*k1*k1;
    }
  // for(int i=0;i<dimk;i++)
  //   {
  //     for(int j=0;j<dimk;j++)
  // 	{
  // 	  h(i,j)+=getVkkfromHF(channel,i,j);
  // 	}
  //   }
  // return;
  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  double lengthb=pSystem_Tz->lengthb;
  int totalGroups=pSystem_Tz->Groups.size();
  int group=0;//find the group of l,jj,tz
  for(group=0;group<totalGroups;group++)
    {
      const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][0] ];
      if(HOOrbital.l==l && HOOrbital.jj==jj && HOOrbital.tz==tz)
	break;
    }
  if(group>=totalGroups) return;
  int dimho=pSystem_Tz->Groups[group].size();
  KMat T(dimk,dimho);
  KMat UHF(dimho,dimho);
  for(int i=0;i<dimk;i++)
    for(int j=0;j<dimho;j++)
      {
	const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][j] ];
	int n=HOOrbital.n;
	KDataType k1=Channels[channel].k[i];
	KDataType wk1=Channels[channel].wk[i];
	T(i,j)=	cmplx_ho_k(n,l,lengthb,k1)*k1 * sqrt(wk1);
      }
  for(int i=0;i<dimho;i++)
    for(int j=i;j<dimho;j++)
      {
	int bra=pSystem_Tz->Groups[group][i];
	int ket=pSystem_Tz->Groups[group][j];
	
        UHF(i,j)=hfsolver.SHFDiagram(bra,ket);
	UHF(j,i)=UHF(i,j);
      }
  h+=T*UHF*T.transpose();
}
void SPSolver::setVkkMat(int channel,KMat &h)
{
  setVkkMatfromHF(channel,h);
  //  setVkkMatfromGF(channel,h);
}
void SPSolver::solve()
{
  Orbitals.clear();
  for(int channel=0;channel<Channels.size();channel++)
    {
      int dim=Channels[channel].numPoints;
      KMat h(dim,dim);
      setVkkMat(channel,h);
      EigenSolver es(h);
      for(int a=0;a<dim;a++)
	{
	  KDataType norm=sqrt( es.eigenvectors().col(a).transpose()*es.eigenvectors().col(a) );
	  if(complexd(100.,0) < es.eigenvalues()[a])
	    continue;
	  Orbitals.push_back( SP_K(Channels[channel].l,Channels[channel].jj,Channels[channel].tz,channel,es.eigenvalues()[a],es.eigenvectors().col(a)/norm) );
	}
    }
  //  setType();
  setupOverlap(Orbitals.size());//set overlap between HO. and gammow basis, used for interaction basis transform

  //the rest partial wave
  int totalGroups=SystemHF.Groups.size();
  for(int group=0;group<totalGroups;group++)
    {
      const System_Tz_HF::OrbitalType & Orb=SystemHF.Orbitals[ SystemHF.Groups[group][0] ];
      int l=Orb.l;
      int jj=Orb.jj;
      int tz=Orb.tz;
      int i=0;
      for(i=0;i<Channels.size();i++)
	{
	  if( Channels[i].l==l && Channels[i].jj==jj && Channels[i].tz==tz )
	    break;
	}
      if(i>=Channels.size())
	{
	  for(int k=0;k<SystemHF.Groups[group].size();k++)
	    {
	      const System_Tz_HF::OrbitalType & HFOrb=SystemHF.Orbitals[ SystemHF.Groups[group][k] ];
	      KDataType e=HFOrb.e;
	      int dim=HFOrb.Coeff.size();
	      SP_K SP(l,jj,tz,-1,e,KVec());
	      SP.GroupIndex=HFOrb.GroupIndex;
	      SP.Coeff.resize(dim);	      
	      for(int n=0;n<dim;n++)
		SP.Coeff[n]=HFOrb.Coeff[n];
	      Orbitals.push_back(SP);
	    }
	}
    }
  stable_sort( Orbitals.begin(),Orbitals.end() );
  
  Energy=0;
  FermiSurface=SystemHF.FermiSurface;
  for(int i=0;i<=FermiSurface;i++)
    {
      KDataType t=0;
      int channel=Orbitals[i].channel;
      if(channel==-1)//HF basis
	{
	  for(int ibra=0;ibra<Orbitals[i].Coeff.size();ibra++)
	    for(int iket=0;iket<Orbitals[i].Coeff.size();iket++)
	      {
		int iibra=pSystem_Tz->Groups[Orbitals[i].GroupIndex][ibra];
		int iiket=pSystem_Tz->Groups[Orbitals[i].GroupIndex][iket];
		t+=pSystem_Tz->get1B(realA,iibra,iiket)*Orbitals[i].Coeff[ibra]*Orbitals[i].Coeff[iket];
	      }	  
	}
      else
	{
	  for(int j=0;j<Channels[channel].numPoints;j++)
	    {
	      KDataType k=Channels[channel].k[j];
	      t+=Orbitals[i].wavefun(j)*(1-1./A)*hbar_c*hbar_c*0.5/mc2*k*k*Orbitals[i].wavefun(j);
	    }
	}
      Energy+=0.5*(Orbitals[i].e+t)*(Orbitals[i].jj+1.);
    }
}

void SPSolver::printOrbitals(int num,int channel)
{
  if(num==-1) num=Orbitals.size();
  int printnum=0;
  for(int i=0;i<Orbitals.size();i++)
    {
      int chan=Orbitals[i].channel;
      if(channel==-1||channel==chan)
	{
	  //	  cout<<"=============================\n";
	  cout<<setw(4)<<Orbitals[i].l<<"\t"<<setw(4)<<Orbitals[i].jj<<"\t"<<setw(4)<<Orbitals[i].tz<<"\t"<<setw(10)<<Orbitals[i].e.real()<<"\t"<<setw(10)<<Orbitals[i].e.imag()<<"\t"<<setw(10)<<std::arg(Orbitals[i].e)*180/Pi<<endl;//<<"\t"<<Orbitals[i].type<<endl;
	  // cout<<"-----------------------------\n";
	  // for(int j=0;j<Channels[chan].numPoints;j++)
	  //   {
	  //     cout<<Channels[chan].k[j].real()<<"\t"<<Channels[chan].wk[j].real()<<"\t"<<Orbitals[i].wavefun(j).real()<<endl;
	  //   }
	  // cout<<"=============================\n\n\n";
	  ++printnum;
	  if(printnum>=num) break;
	}
    }
}

void SPSolver::printVkk(int channel)
{
  int dim=Channels[channel].numPoints;
  for(int a=0;a<dim;a++)
    for(int b=0;b<dim;b++)
      {
	KDataType k1=Channels[channel].k[a];
	KDataType k2=Channels[channel].k[b];
	cout<<k1<<"\t"<<k2<<"\t"<<getVkkfromHF(channel,a,b)<<endl;
      }
}

void SPSolver::setType()
{
  int index=0;
  for(int channel=0;channel<Channels.size();channel++)
    {
      int dim=Channels[channel].numPoints;
      int linedim=dim/2;
      vector<double> thetas(linedim);
      for(int i=0;i<dim;i++)
	{
	  Orbitals[index+i].type=1;
	  if(i<linedim)
	    {
	      if(Orbitals[index+i].e.real()<0)
		Orbitals[index+i].type=-1;
	      else
		{
		  thetas[i]=arg(Orbitals[index+i].e);
		}
	    }
	}
      double sum = accumulate(thetas.begin(), thetas.end(), 0.0);
      double mean = sum / linedim;
      for(int i=0;i<linedim;i++)
	{
	  //	  cout<<thetas[i]-mean<<endl;
	  if(abs(thetas[i]-mean)>0.1)
	    Orbitals[index+i].type=0;
	}
      index+=dim;
    }
}


void SPSolver::setupOverlap(int num)
{
  double lengthb=pSystem_Tz->lengthb;
  int totalGroups=pSystem_Tz->Groups.size();

  for(int i=0;i<num;i++)
    {
      int channel=Orbitals[i].channel;
      int l=Channels[channel].l;
      int jj=Channels[channel].jj;
      int tz=Channels[channel].tz;
      int group=0;//find the group of l,jj,tz
      for(group=0;group<totalGroups;group++)
	{
	  const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][0] ];
	  if(HOOrbital.l==l && HOOrbital.jj==jj && HOOrbital.tz==tz)
	    break;
	}
      if(group>=totalGroups)
	{
	  std::cerr<<"there are channels in spsolver orbitals not found in HO. orbitals:\n";
	  std::cerr<<l<<"\t"<<jj<<"\t"<<tz<<std::endl;
	  exit(-1);
	}
      
      Orbitals[i].GroupIndex=group;
      int dimho=pSystem_Tz->Groups[group].size();

      Orbitals[i].Coeff.resize(dimho);
      for(int j=0;j<dimho;j++)//every HO. orbitals in a group
	{
	  const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][j] ];
	  int n=HOOrbital.n;	  
	  KDataType temp=0;//overlap with the HO. orbital
	  for(int k=0;k<Channels[channel].numPoints;k++)
	    {
	      KDataType k1=Channels[channel].k[k];
	      KDataType wk1=Channels[channel].wk[k];
	      temp += cmplx_ho_k(n,l,lengthb,k1)*k1*sqrt(wk1)*Orbitals[i].wavefun(k);
	    }
	  Orbitals[i].Coeff[j]=temp;
	  //	  std::cout<<temp<<std::endl;
	}
    }
}
